from nicegui import ui
import pandas as pd
import plotly.express as px

# 载入数据
df = pd.read_csv('personality_datasert.csv')

# 将英文性格类型转换为中文
personality_mapping = {'Introvert': '内向', 'Extrovert': '外向'}
df['Personality'] = df['Personality'].map(personality_mapping)

# 提取唯一性格类别
personality_types = df['Personality'].unique()

# 过滤数据的函数
def filter_data(personality_filter, min_time_alone, max_time_alone):
    filtered = df.copy()
    if personality_filter != '全部':
        filtered = filtered[filtered['Personality'] == personality_filter]
    filtered = filtered[
        (filtered['Time_spent_Alone'] >= min_time_alone) &
        (filtered['Time_spent_Alone'] <= max_time_alone)
        ]
    return filtered

# 获取性格类型统计
def get_personality_counts(data_frame):
    return data_frame['Personality'].value_counts()

# 创建初始图表
def create_initial_charts():
    # 饼图
    initial_counts = get_personality_counts(df)
    pie_fig = px.pie(
        names=initial_counts.index,
        values=initial_counts.values,
        title='性格分布',
        color_discrete_map={'内向': 'blue', '外向': 'orange'}
    )
    pie_fig.update_traces(textinfo='percent+label')

    # 3D散点图
    scatter3d_fig = px.scatter_3d(
        df,
        x='Time_spent_Alone',
        y='Social_event_attendance',
        z='Going_outside',
        color='Personality',
        title='3D散点图：独处时间vs社交活动vs外出频率',
        labels={
            'Time_spent_Alone': '独处时间（小时）',
            'Social_event_attendance': '社交活动参与度',
            'Going_outside': '外出频率',
            'Personality': '性格类型'
        },
        color_discrete_map={'内向': 'blue', '外向': 'orange'}
    )
    scatter3d_fig.update_traces(marker=dict(size=5))

    return pie_fig, scatter3d_fig

# 创建初始图表
pie_fig, scatter3d_fig = create_initial_charts()

# UI布局
with ui.row().classes('p-4'):
    # 左侧筛选面板
    with ui.column().classes('w-1/4'):
        ui.label('筛选条件').classes('text-xl font-bold mb-2')

        personality_select = ui.select(
            label='性格类型',
            options=['全部'] + list(personality_types),
            value='全部'
        )

        ui.label('最少独处时间（小时）')
        min_time_slider = ui.slider(min=0, max=15, value=0, step=0.5)

        ui.label('最多独处时间（小时）')
        max_time_slider = ui.slider(min=0, max=15, value=15, step=0.5)

        # 应用按钮
        def update():
            filtered = filter_data(
                personality_select.value,
                min_time_slider.value,
                max_time_slider.value
            )
            update_charts(filtered)

        ui.button('应用筛选', on_click=update).classes('mt-2')

    # 右侧图表区
    with ui.column().classes('w-3/4'):
        # 饼图
        ui.label('性格类别分布').classes('text-xl font-bold')
        pie_chart = ui.plotly(pie_fig).style('height: 400px')

        # 3D散点图
        ui.label('3D散点图').classes('text-xl font-bold mt-4')
        scatter3d_chart = ui.plotly(scatter3d_fig).style('height: 600px')

# 更新图表函数
def update_charts(filtered_df):
    # 更新饼图
    filtered_counts = get_personality_counts(filtered_df)
    if len(filtered_counts) > 0:
        updated_pie_fig = px.pie(
            names=filtered_counts.index,
            values=filtered_counts.values,
            title='性格类别分布（已筛选）',
            color_discrete_map={'内向': 'blue', '外向': 'orange'}
        )
        updated_pie_fig.update_traces(textinfo='percent+label')
        pie_chart.update_figure(updated_pie_fig)

    # 更新3D散点图
    if len(filtered_df) > 0:
        updated_scatter3d_fig = px.scatter_3d(
            filtered_df,
            x='Time_spent_Alone',
            y='Social_event_attendance',
            z='Going_outside',
            color='Personality',
            title='3D散点图：独处时间vs社交活动vs外出频率（已筛选）',
            labels={
                'Time_spent_Alone': '独处时间（小时）',
                'Social_event_attendance': '社交活动参与度',
                'Going_outside': '外出频率',
                'Personality': '性格类型'
            },
            color_discrete_map={'内向': 'blue', '外向': 'orange'}
        )
        updated_scatter3d_fig.update_traces(marker=dict(size=5))
        scatter3d_chart.update_figure(updated_scatter3d_fig)

# 使用兼容多进程的主程序保护
# if __name__ in {"__main__", "__mp_main__"}:
#     ui.run()
    if __name__ in {"__main__", "__mp_main__"}:
        import sys

        port = 8081
        if '--port' in sys.argv:
            port_index = sys.argv.index('--port')
            if port_index + 1 < len(sys.argv):
                port = int(sys.argv[port_index + 1])
        ui.run(title='基础性格分析', port=port, show=True)